#!/usr/bin/env python

import numpy as np
import math
from scipy.stats import *



def count_occurences(x,y=None):
    """
    count_occurencens(): Return unique values in x and number
                         of times each value appears. If provided
                         with an additional array, y, will count
                         the numberof times each value of y appears
                         in x.

    USAGE: u,n = count_occurences(x,y)

           x  = array of integers
           y  = array of unique values to look for in x
           u = array of unique values in x or y
           n = number of times each value in u appears in x
    """
    if y is None:
        ux,inv = np.unique(x,return_inverse=True)
        nx = np.bincount(inv)
        return ux,nx
    else:
        isort = np.argsort(y)
        jsort = np.argsort(isort)
        counts,extent = np.histogram(x,np.append(np.copy(y[isort]),\
                                                 np.max(y)+1))
        return y,counts[jsort]


def histogram_data(x,xbins,weights=None,cumulative=False):
    bw = xbins[1] - xbins[0]
    bins = np.copy(xbins)    
    bins = np.append(bins,np.ones(1)*(xbins[-1]+bw))
    if weights is None:
        weights = np.ones_like(x)
        if not cumulative:
            weights = weights*(1.0/bw)
    hist,dummy = np.histogram(x,bins,weights=weights)
    if cumulative:
        hist = np.cumsum(hist)
    return hist

        
    
def random_sample_mask(n,percent):
    """
    random_sample(): Returns a mask array with single
                     dimension 'n', that will select
                     approximately the specified
                     percentage of data (assuming data
                     is in an array also of size 'n')
    USAGE: mask = random_sample(n,percent)

           n       : length of mask array
           percent : percent to select
           mask    : logical array with specified
                     percentage of values set to True
                      
    """
    import random
    frac = float(percent)/100.0
    frac = int(np.ceil(frac*float(n)))
    index = random.sample(np.arange(int(n)),frac)
    mask = np.zeros(n,bool)
    mask[index] = True
    return mask



def binstats(X,Y,Xbins,statistic="median",weights=None,percentile=68):
    statistic = str(statistic)
    if statistic.lower().startswith("perc"):
        def perc(x):
            return np.percentile(x,q=percentile)
        statistic = perc
    elif statistic.lower() in ["average","avg"]:
        def avg(x):
            return np.average(x,weights=weights)
        statistic = avg
    elif statistic.lower() == "std":
        statistic = np.std
    elif statistic.lower() == "var":
        statistic = np.var
    elif statistic.lower().startswith("prod"):
        statistic = np.prod
    elif statistic.lower().startswith("min"):
        statistic = np.nanmin
    elif statistic.lower().startswith("max"):
        statistic = np.nanmax
    elif statistic.lower() == "mode":
        def mode_of_data(x):
            return mode(x)[0][0]
        statistic = mode_of_data
    else:
        pass
    dX = np.array([Xbins[-1]]) + (Xbins[1] - Xbins[0])
    Xbins = np.append(Xbins,dX)
    result,bin_edges,binnumber = binned_statistic(X,Y,statistic=statistic,bins=Xbins)
    return result,bin_edges,binnumber 
    


def binstats2D(X,Y,Z,Xbins,Ybins,statistic="median",weights=None,percentile=68):
    """
    statistic can be: mean,median,sum,product,std,var,percentile,avg,max,min,mode

    NB 'avg' is weighted average
    """
    bins = [Xbins,Ybins]
    statistic = str(statistic)
    if statistic.lower().startswith("perc"):
        def perc(x):
            return np.percentile(x,q=percentile)
        statistic = perc
    elif statistic.lower() in ["average","avg"]:
        def avg(x):
            return np.average(x,weights=weights)
        statistic = avg
    elif statistic.lower() == "std":
        statistic = np.std
    elif statistic.lower() == "var":
        statistic = np.var
    elif statistic.lower().startswith("prod"):
        statistic = np.prod
    elif statistic.lower().startswith("min"):
        statistic = np.nanmin
    elif statistic.lower().startswith("max"):
        statistic = np.nanmax
    elif statistic.lower() == "mode":
        def mode_of_data(x):
            return mode(x)[0][0]
        statistic = mode_of_data
    else:
        pass
    stat,xedges,yedges,numb = binned_statistic_2d(X,Y,Z,statistic=statistic,bins=bins)
    return stat,xedges,yedges,numb



def Inf2NaN2D(arr):
    """
    Inf2NaN2d(): Converts all +/-Inf in a 2D array to NaN.
                 (Useful when plotting images of arrays that
                 have been divided or logged).

    USAGE: arr = Inf2NaN2D(arr)

           arr : input/output array, with +/Inf in input
                 and NaN on output.
                 
    """
    mask = np.isinf(arr)
    arr[mask] = np.NaN
    return arr


def AND(*args):
    return [all(tuple) for tuple in zip(*args)]

def OR(*args):
    return [any(tuple) for tuple in zip(*args)]

    
def chi_squared(ex,ob,er):
    chi_sq = np.sum(((ob-ex)/er)**2)
    return chi_sq



def matchIDs(arr1,arr2,arr2_sorted=False):
    """
    For each element in arr1 return the index of the element with the
    same value in arr2, or -1 if there is no element with the same value.
    Setting arr2_sorted=True will save some time if arr2 is already sorted
    into ascending order.
    
    It is assumed that each element in arr1 only occurs once in arr2.
    """    
    # Workaround for a numpy bug (<=1.4): ensure arrays are native endian
    # because searchsorted ignores endian flag
    if not(arr1.dtype.isnative):
        arr1_n = asarray(arr1, dtype=arr1.dtype.newbyteorder("="))
    else:
        arr1_n = arr1
    if not(arr2.dtype.isnative):
        arr2_n = asarray(arr2, dtype=arr2.dtype.newbyteorder("="))
    else:
        arr2_n = arr2
    # Sort arr2 into ascending order if necessary
    tmp1 = arr1_n
    if arr2_sorted:
        tmp2 = arr2_n
        idx = slice(0,len(arr2_n))
    else:
        idx = argsort(arr2_n)
        tmp2 = arr2_n[idx]
    # Find where elements of arr1 are in arr2
    ptr  = searchsorted(tmp2, tmp1)
    # Make sure all elements in ptr are valid indexes into tmp2
    # (any out of range entries won't match so they'll get set to -1
    # in the next bit)
    ptr[ptr>=len(tmp2)] = 0
    ptr[ptr<0]          = 0
    # Return -1 where no match is found                                                                                                                                                                            
    ind  = tmp2[ptr] != tmp1
    ptr[ind] = -1
    # Put ptr back into original order
    ind = arange(len(arr2_n))[idx]
    ptr = where(ptr>= 0, ind[ptr], -1)
    # Return pointer
    return ptr
                    
